from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain.agents import create_structured_chat_agent, AgentExecutor
from langchain.memory import ConversationBufferMemory
from langchain.tools import Tool
from langchain_experimental.tools import PythonREPLTool
from langchain_experimental.agents.agent_toolkits import create_python_agent

key = 'key'
model = ChatOpenAI(model="gpt-3.5-turbo",
                   openai_api_key=key,
                   openai_api_base="https://api.aigc369.com/v1",
                   temperature=0)

agent_python_executor = create_python_agent(
    llm=model,
    tool=PythonREPLTool(),
    verbose=True,
    agent_executor_kwargs={
        "handle_parsing_errors": True
    }
)


@tool
def get_text_len(text: str) -> int:
    """这个工具是用来获取文本长度的，如果你有需要进行文本长度获取的任务，可以使用这个工具"""
    return len(text)


tools = [get_text_len,
         Tool(
             name="python解释器",
             description="""当你需要借助Python解释器时，比如进行数学计算等操作时使用这个工具。
                用自然语言把要求或者要求中的表达式提取给这个工具，生成要求对应的Python代码并返回代码执行的结果。""",
             func=agent_python_executor
         )
         ]

react_prompt_content = """
    Respond to the human as helpfully and accurately as possible. You have access to the following tools:
    
    {tools}
    
    Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).
    
    Valid "action" values: "Final Answer" or {tool_names}
    
    Provide only ONE action per $JSON_BLOB, as shown:
    
    ```
    {{
      "action": $TOOL_NAME,
      "action_input": $INPUT
    }}
    ```
    
    Follow this format:
    
    Question: input question to answer
    Thought: consider previous and subsequent steps
    Action:
    ```
    $JSON_BLOB
    ```
    Observation: action result
    ... (repeat Thought/Action/Observation N times)
    Thought: I know what to respond
    Action:
    ```
    {{
      "action": "Final Answer",
      "action_input": "Final response to human"
    }}
    
    Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation"),
      ("placeholder", "{chat_history}"),
      ("human", "{input}
    
    {agent_scratchpad}
     (reminder to respond in a JSON blob no matter what)")
 """

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", react_prompt_content),
        ("human", "{input}")
    ]
)

agent = create_structured_chat_agent(
    llm=model,
    tools=tools,
    prompt=prompt
)

memory = ConversationBufferMemory(
    return_history=True,
    memory_key="chat_history"
)

agent_executor = AgentExecutor.from_agent_and_tools(
    tools=tools,
    agent=agent,
    memory=memory,
    verbose=True,
    handle_parsing_errors=True
)

# print(agent_executor.invoke({"input": "计算一下3* 5等于多少"}))
print(agent_executor.invoke({"input": "我叫aa"}))
print(agent_executor.invoke({"input": "我是谁"}))
# print(agent_executor.invoke({"input": "君不见黄河之水天上来，奔流到海不复回'。这句话的长度是多少?"}))


# model_with_tools = model.bind_tools(tools)
# response = model_with_tools.invoke("君不见黄河之水天上来，奔流到海不复回'。这句话的长度是多少?").tool_calls
#
# # print(response)
# #
# # _input = response[0]['args']
# #
# # print(get_text_len.invoke(_input))
#
# chain = model_with_tools | (lambda x: x.tool_calls[0]['args']) | get_text_len
#
# print(chain.invoke("君不见黄河之水天上来，奔流到海不复回'。这句话的长度是多少?"))
